#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pick_fastest_pip.py
自动测速常见 PyPI 镜像（pip 源），选择最快者，并可持久化为当前用户默认源。

设计要点：
- 仅用标准库（urllib、ssl、time、subprocess、argparse），零外部依赖。
- 通过读取每个候选源的 /simple/<test-package>/ 页面，测首字节延迟 + 读取速度。
- 评估指标：综合耗时（连接 + 首包 + 连续读取上限 bytes），取多次尝试的最优值。
- 持久化：调用 `python -m pip config set global.index-url <url>`，不修改系统级。
  如需回滚：python -m pip config unset global.index-url
"""

import argparse
import sys
import time
import urllib.request
import urllib.parse
import ssl
import subprocess
from typing import List, Tuple, Optional

DEFAULT_CANDIDATES = [
    # 官方 PyPI
    "https://pypi.org/simple",
    # 常见国内镜像（可按需增删）
    "https://pypi.tuna.tsinghua.edu.cn/simple",
    "https://mirrors.aliyun.com/pypi/simple",
    "https://repo.huaweicloud.com/repository/pypi/simple",
    "https://mirrors.cloud.tencent.com/pypi/simple",
    # USTC 近年不稳定，保留为候选但失败会自动跳过
    "https://pypi.mirrors.ustc.edu.cn/simple",
]

UA = "pip-source-bench/1.0 (+https://python.org)"


def format_secs(s: float) -> str:
    return f"{s*1000:.0f} ms" if s < 1 else f"{s:.2f} s"


def bench_one(index_url: str, test_package: str, max_bytes: int, timeout: float, tries: int
              ) -> Tuple[str, Optional[float], str]:
    """
    基于 /simple/<test_package>/ 页面进行读取测速。
    返回：(index_url, 最佳耗时(秒) 或 None, 附加说明)
    """
    url = index_url.rstrip("/") + f"/{test_package.strip('/')}/"
    best = None
    note = ""
    ctx = ssl.create_default_context()

    for i in range(tries):
        start = time.perf_counter()
        first_byte_at = None
        read_bytes = 0
        try:
            req = urllib.request.Request(url, headers={"User-Agent": UA})
            with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
                # 记录首字节到达时间
                chunk = resp.read(1)
                if not chunk:
                    # 空响应，直接判失败
                    note = "empty response"
                    continue
                first_byte_at = time.perf_counter()
                read_bytes += len(chunk)

                # 继续读取直到 max_bytes 或响应结束
                chunk_size = 8192
                while read_bytes < max_bytes:
                    chunk = resp.read(min(chunk_size, max_bytes - read_bytes))
                    if not chunk:
                        break
                    read_bytes += len(chunk)

            elapsed = time.perf_counter() - start
            # 采用总耗时（连接+首字节+读取）的最小值作为该源的分数
            if best is None or elapsed < best:
                best = elapsed
        except Exception as e:
            note = f"error: {type(e).__name__}: {e}"
            # 尝试下一个重试
            continue

    if best is None:
        return index_url, None, note or "unreachable"

    extra = ""
    if first_byte_at is not None:
        rtt_like = first_byte_at - start
        extra = f" | RTT≈{format_secs(rtt_like)} | read={read_bytes} B"
    return index_url, best, f"ok{extra}"


def rank(candidates: List[str], test_package: str, max_bytes: int, timeout: float, tries: int):
    results = []
    for url in candidates:
        url = url.strip()
        if not url:
            continue
        idx, best, note = bench_one(url, test_package, max_bytes, timeout, tries)
        results.append((idx, best, note))

    # 按耗时升序；不可达放最后
    results.sort(key=lambda x: float("inf") if x[1] is None else x[1])
    return results


def apply_index_url(index_url: str) -> Tuple[bool, str]:
    """
    通过 `python -m pip config set global.index-url <url>` 写入用户级配置。
    返回 (成功与否, 输出/错误信息)
    """
    cmd = [sys.executable, "-m", "pip", "config", "set", "global.index-url", index_url]
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, check=False)
        ok = (p.returncode == 0)
        out = (p.stdout or "") + (p.stderr or "")
        return ok, out.strip()
    except Exception as e:
        return False, f"spawn error: {type(e).__name__}: {e}"


def show_current_config() -> str:
    cmd = [sys.executable, "-m", "pip", "config", "list"]
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, check=False)
        return (p.stdout or "") + (p.stderr or "")
    except Exception as e:
        return f"cannot read config: {e}"


def main():
    ap = argparse.ArgumentParser(
        description="自动测速并设置最快的 pip 源（PyPI 镜像）。"
    )
    ap.add_argument("--candidates", type=str, default="",
                    help="自定义候选源，逗号分隔（会与内置列表合并，重复自动去重）。")
    ap.add_argument("--test-package", type=str, default="pip",
                    help="用于测试的包名（默认：pip，对应 /simple/pip/ 页面）。")
    ap.add_argument("--bytes", type=int, default=512 * 1024,
                    help="单源最大读取字节上限（默认：524288 字节=512KiB）。")
    ap.add_argument("--timeout", type=float, default=10.0,
                    help="单次连接+读的超时（秒，默认：10.0）。")
    ap.add_argument("--tries", type=int, default=2,
                    help="每个源的尝试次数（默认：2，取最好的一次）。")
    ap.add_argument("--apply", action="store_true",
                    help="将最快源持久化为当前用户默认源（使用 pip config set）。")
    args = ap.parse_args()

    # 合并候选
    cand = list(DEFAULT_CANDIDATES)
    if args.candidates.strip():
        cand += [x.strip() for x in args.candidates.split(",") if x.strip()]
    # 去重同时保持顺序
    seen = set()
    uniq = []
    for u in cand:
        if u not in seen:
            uniq.append(u)
            seen.add(u)

    print("== pip 源测速开始 ==")
    print(f"候选数量：{len(uniq)} | 测试包：{args.test_package} | 读取上限：{args.bytes} B | 重试：{args.tries}")
    results = rank(uniq, args.test_package, args.bytes, args.timeout, args.tries)

    print("\n== 速度排名（越上越快）==")
    best_ok = None
    for i, (u, t, note) in enumerate(results, 1):
        if t is None:
            speed = "不可达/失败"
        else:
            speed = format_secs(t)
            if best_ok is None:
                best_ok = u
        print(f"{i:02d}. {u}  ->  {speed}   [{note}]")

    if best_ok is None:
        print("\n未找到可用源。请检查网络、代理（proxy）或更换候选列表。")
        sys.exit(2)

    print(f"\n最快候选：{best_ok}")

    # 显示当前 pip 配置
    print("\n== 当前 pip 配置（节选）==")
    print(show_current_config())

    if args.apply:
        print("\n== 写入默认源 ==")
        ok, out = apply_index_url(best_ok)
        if ok:
            print(f"已将默认源设置为：{best_ok}")
        else:
            print("写入失败：")
        print(out)

        print("\n== 更新后 pip 配置（节选）==")
        print(show_current_config())
    else:
        print("\n（仅测速未写入。如需持久化，追加参数：--apply）")


if __name__ == "__main__":
    main()
